import paddle
from utils.activations import h_sigmoid


class SELayer(paddle.nn.Layer):

    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = paddle.nn.AdaptiveAvgPool2D(output_size=1)
        self.conv1 = paddle.nn.Conv2D(in_channels=channel, out_channels=
            channel // reduction, kernel_size=1, padding=0)
        self.relu = paddle.nn.ReLU()
        self.conv2 = paddle.nn.Conv2D(in_channels=channel // reduction,
            out_channels=channel, kernel_size=1, padding=0)
        self.sigmoid = h_sigmoid()

    def forward(self, x):
        module_input = x
        x = self.avg_pool(x)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.sigmoid(x)
        return module_input * x
